{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7765UFHoyGx6"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KsOkK8O69PyT"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZS8z-_KeywY9"
      },
      "source": [
        "# TF Lattice Custom Estimators"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r61fkA2i9Y3_"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/lattice/tutorials/custom_estimators\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/lattice/blob/master/docs/tutorials/custom_estimators.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/lattice/blob/master/docs/tutorials/custom_estimators.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/lattice/docs/tutorials/custom_estimators.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ur6yCw7YVvr8"
      },
      "source": [
        "## Overview\n",
        "\n",
        "You can use custom estimators to create arbitrarily monotonic models using TFL layers. This guide outlines the steps needed to create such estimators."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x769lI12IZXB"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fbBVAR6UeRN5"
      },
      "source": [
        "Installing TF Lattice package:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bpXjJKpSd3j4"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "!pip install tensorflow-lattice"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jSVl9SHTeSGX"
      },
      "source": [
        "Importing required packages:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "P9rMpg1-ASY3"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "import logging\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import sys\n",
        "import tensorflow_lattice as tfl\n",
        "from tensorflow import feature_column as fc\n",
        "\n",
        "from tensorflow_estimator.python.estimator.canned import optimizers\n",
        "from tensorflow_estimator.python.estimator.head import binary_class_head\n",
        "logging.disable(sys.maxsize)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "svPuM6QNxlrH"
      },
      "source": [
        "Downloading the UCI Statlog (Heart) dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "M0CmH1gPASZF"
      },
      "outputs": [],
      "source": [
        "csv_file = tf.keras.utils.get_file(\n",
        "    'heart.csv', 'http://storage.googleapis.com/applied-dl/heart.csv')\n",
        "df = pd.read_csv(csv_file)\n",
        "target = df.pop('target')\n",
        "train_size = int(len(df) * 0.8)\n",
        "train_x = df[:train_size]\n",
        "train_y = target[:train_size]\n",
        "test_x = df[train_size:]\n",
        "test_y = target[train_size:]\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nKkAw12SxvGG"
      },
      "source": [
        "Setting the default values used for training in this guide:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "1T6GFI9F6mcG"
      },
      "outputs": [],
      "source": [
        "LEARNING_RATE = 0.1\n",
        "BATCH_SIZE = 128\n",
        "NUM_EPOCHS = 1000"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0TGfzhPHzpix"
      },
      "source": [
        "## Feature Columns\n",
        "\n",
        "As for any other TF estimator, data needs to be passed to the estimator, which is typically via an input_fn and parsed using [FeatureColumns](https://www.tensorflow.org/guide/feature_columns)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DCIUz8apzs0l"
      },
      "outputs": [],
      "source": [
        "# Feature columns.\n",
        "# - age\n",
        "# - sex\n",
        "# - ca        number of major vessels (0-3) colored by flourosopy\n",
        "# - thal      3 = normal; 6 = fixed defect; 7 = reversable defect\n",
        "feature_columns = [\n",
        "    fc.numeric_column('age', default_value=-1),\n",
        "    fc.categorical_column_with_vocabulary_list('sex', [0, 1]),\n",
        "    fc.numeric_column('ca'),\n",
        "    fc.categorical_column_with_vocabulary_list(\n",
        "        'thal', ['normal', 'fixed', 'reversible']),\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hEZstmtT2CA3"
      },
      "source": [
        "Note that categorical features do not need to be wrapped by a dense feature column, since `tfl.laysers.CategoricalCalibration` layer can directly consume category indices."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H_LoW_9m5OFL"
      },
      "source": [
        "## Creating input_fn\n",
        "\n",
        "As for any other estimator, you can use input_fn to feed data to the model for training and evaluation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lFVy1Efy5NKD"
      },
      "outputs": [],
      "source": [
        "train_input_fn = tf.compat.v1.estimator.inputs.pandas_input_fn(\n",
        "    x=train_x,\n",
        "    y=train_y,\n",
        "    shuffle=True,\n",
        "    batch_size=BATCH_SIZE,\n",
        "    num_epochs=NUM_EPOCHS,\n",
        "    num_threads=1)\n",
        "\n",
        "test_input_fn = tf.compat.v1.estimator.inputs.pandas_input_fn(\n",
        "    x=test_x,\n",
        "    y=test_y,\n",
        "    shuffle=False,\n",
        "    batch_size=BATCH_SIZE,\n",
        "    num_epochs=1,\n",
        "    num_threads=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kbrgSr9KaRg0"
      },
      "source": [
        "## Creating model_fn\n",
        "\n",
        "There are several ways to create a custom estimator. Here we will construct a `model_fn` that calls a Keras model on the parsed input tensors. To parse the input features, you can use `tf.feature_column.input_layer`, `tf.keras.layers.DenseFeatures`, or `tfl.estimators.transform_features`. If you use the latter, you will not need to wrap categorical features with dense feature columns, and the resulting tensors will not be concatenated, which makes it easier to use the features in the calibration layers.\n",
        "\n",
        "To construct a model, you can mix and match TFL layers or any other Keras layers. Here we create a calibrated lattice Keras model out of TFL layers and impose several monotonicity constraints. We then use the Keras model to create the custom estimator.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n2Zrv6OPaQO2"
      },
      "outputs": [],
      "source": [
        "def model_fn(features, labels, mode, config):\n",
        "  \"\"\"model_fn for the custom estimator.\"\"\"\n",
        "  del config\n",
        "  input_tensors = tfl.estimators.transform_features(features, feature_columns)\n",
        "  inputs = {\n",
        "      key: tf.keras.layers.Input(shape=(1,), name=key) for key in input_tensors\n",
        "  }\n",
        "\n",
        "  lattice_sizes = [3, 2, 2, 2]\n",
        "  lattice_monotonicities = ['increasing', 'none', 'increasing', 'increasing']\n",
        "  lattice_input = tf.keras.layers.Concatenate(axis=1)([\n",
        "      tfl.layers.PWLCalibration(\n",
        "          input_keypoints=np.linspace(10, 100, num=8, dtype=np.float32),\n",
        "          # The output range of the calibrator should be the input range of\n",
        "          # the following lattice dimension.\n",
        "          output_min=0.0,\n",
        "          output_max=lattice_sizes[0] - 1.0,\n",
        "          monotonicity='increasing',\n",
        "      )(inputs['age']),\n",
        "      tfl.layers.CategoricalCalibration(\n",
        "          # Number of categories including any missing/default category.\n",
        "          num_buckets=2,\n",
        "          output_min=0.0,\n",
        "          output_max=lattice_sizes[1] - 1.0,\n",
        "      )(inputs['sex']),\n",
        "      tfl.layers.PWLCalibration(\n",
        "          input_keypoints=[0.0, 1.0, 2.0, 3.0],\n",
        "          output_min=0.0,\n",
        "          output_max=lattice_sizes[0] - 1.0,\n",
        "          # You can specify TFL regularizers as tuple\n",
        "          # ('regularizer name', l1, l2).\n",
        "          kernel_regularizer=('hessian', 0.0, 1e-4),\n",
        "          monotonicity='increasing',\n",
        "      )(inputs['ca']),\n",
        "      tfl.layers.CategoricalCalibration(\n",
        "          num_buckets=3,\n",
        "          output_min=0.0,\n",
        "          output_max=lattice_sizes[1] - 1.0,\n",
        "          # Categorical monotonicity can be partial order.\n",
        "          # (i, j) indicates that we must have output(i) <= output(j).\n",
        "          # Make sure to set the lattice monotonicity to 'increasing' for this\n",
        "          # dimension.\n",
        "          monotonicities=[(0, 1), (0, 2)],\n",
        "      )(inputs['thal']),\n",
        "  ])\n",
        "  output = tfl.layers.Lattice(\n",
        "      lattice_sizes=lattice_sizes, monotonicities=lattice_monotonicities)(\n",
        "          lattice_input)\n",
        "\n",
        "  training = (mode == tf.estimator.ModeKeys.TRAIN)\n",
        "  model = tf.keras.Model(inputs=inputs, outputs=output)\n",
        "  logits = model(input_tensors, training=training)\n",
        "\n",
        "  if training:\n",
        "    optimizer = optimizers.get_optimizer_instance_v2('Adagrad', LEARNING_RATE)\n",
        "  else:\n",
        "    optimizer = None\n",
        "\n",
        "  head = binary_class_head.BinaryClassHead()\n",
        "  return head.create_estimator_spec(\n",
        "      features=features,\n",
        "      mode=mode,\n",
        "      labels=labels,\n",
        "      optimizer=optimizer,\n",
        "      logits=logits,\n",
        "      trainable_variables=model.trainable_variables,\n",
        "      update_ops=model.updates)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mng-VtsSbVtQ"
      },
      "source": [
        "## Training and Estimator\n",
        "\n",
        "Using the `model_fn` we can create and train the estimator."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j38GaEbKbZju"
      },
      "outputs": [],
      "source": [
        "estimator = tf.estimator.Estimator(model_fn=model_fn)\n",
        "estimator.train(input_fn=train_input_fn)\n",
        "results = estimator.evaluate(input_fn=test_input_fn)\n",
        "print('AUC: {}'.format(results['auc']))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "custom_estimators.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
